package com.ztom.agent.asm.plugin.http.servletrequest;

import com.ztom.agent.common.TraceHeaders;
import com.ztom.agent.trace.Span;
import com.ztom.agent.util.TraceAndSpanIdGenerator;

/**
 * @author ZhangTao
 */
public class HttpRequestTracingUtils {

    public static final String UNSPECIFIED_SPAN_NAME = "UNSPECIFIED";

    private HttpRequestTracingUtils() {
        // Nothing to do
    }

    /**
     * 从请求中获取Span
     *
     * @param request
     * @return
     */
    public static Span fromRequestWithHeaders(RequestWithHeaders request) {
        if (request == null) {
            return null;
        }

        String traceId = getTraceId(request);
        if (traceId == null) {
            return null;
        }

        String spanName = getHeaderWithAttributeAsBackup(request, TraceHeaders.SPAN_NAME);
        if (spanName == null || spanName.length() == 0) {
            spanName = UNSPECIFIED_SPAN_NAME;
        }

        return Span.newBuilder(spanName, Span.SpanType.HTTP_REQUEST)
                .withTraceId(traceId)
                .withParentSpanId(getSpanIdFromRequest(request, TraceHeaders.PARENT_SPAN_ID, false))
                .withSpanId(getSpanIdFromRequest(request, TraceHeaders.SPAN_ID, true))
                .withSampleable(getSpanSampleableFlag(request))
                .build();
    }

    /**
     * 从请求头中获取是否取样的标志
     *
     * @param request
     * @return
     */
    protected static boolean getSpanSampleableFlag(RequestWithHeaders request) {
        String spanSampleableHeaderStr = getHeaderWithAttributeAsBackup(request, TraceHeaders.TRACE_SAMPLED);
        // Default to true (enabling trace sampling for requests that don't explicitly exclude it)
        boolean result = true;

        if ("0".equals(spanSampleableHeaderStr) || "false".equalsIgnoreCase(spanSampleableHeaderStr)) {
            result = false;
        }

        return result;
    }

    /**
     * 从请求头中获取SpanId或ParentSpanId
     *
     * @param request
     * @param headerName
     * @param generateNewSpanIdIfNotFoundInRequest 如果没拿到是否新建SpanId
     * @return
     */
    protected static String getSpanIdFromRequest(RequestWithHeaders request, String headerName, boolean generateNewSpanIdIfNotFoundInRequest) {
        String spanIdString = getHeaderWithAttributeAsBackup(request, headerName);
        if (spanIdString == null) {
            return generateNewSpanIdIfNotFoundInRequest ? TraceAndSpanIdGenerator.generateId() : null;
        }

        return spanIdString;
    }

    /**
     * 从请求头中获取TraceId
     *
     * @param request
     * @return
     */
    protected static String getTraceId(RequestWithHeaders request) {
        return getHeaderWithAttributeAsBackup(request, TraceHeaders.TRACE_ID);
    }

    /**
     * 从请求头中获取headerName的对应值, 如果没取到, 则从Attribute中拿
     *
     * @param request
     * @param headerName
     * @return
     */
    protected static String getHeaderWithAttributeAsBackup(RequestWithHeaders request, String headerName) {
        Object result = request.getHeader(headerName);

        if (result == null || result.toString().trim().length() == 0) {
            result = request.getAttribute(headerName);
        }

        return (result == null) ? null : result.toString().trim();
    }
}
